import sys
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QComboBox, QStyledItemDelegate, QVBoxLayout, QWidget
)
from PySide6.QtCore import Qt
from PySide6.QtGui import QColor, QBrush


class ColoredItemDelegate(QStyledItemDelegate):
    def __init__(self, parent=None):
        super().__init__(parent)
        # 定义四个固定颜色
        self.colors = [
            QColor("#FFCCCB"),  # 浅红色
            QColor("#90EE90"),  # 浅绿色
            QColor("#ADD8E6"),  # 浅蓝色
            QColor("#FFFFE0")   # 浅黄色
        ]
    
    def initStyleOption(self, option, index):
        super().initStyleOption(option, index)
        
        # 根据索引设置背景颜色（只处理前四个选项）
        if 0 <= index.row() < 4:
            option.backgroundBrush = QBrush(self.colors[index.row()])


class ColoredComboBox(QComboBox):
    def __init__(self, items=None, parent=None):
        super().__init__(parent)
        
        # 定义四个固定颜色
        self.colors = [
            "#FFCCCB",  # 浅红色
            "#90EE90",  # 浅绿色
            "#ADD8E6",  # 浅蓝色
            "#FFFFE0"   # 浅黄色
        ]
        
        # 添加四个固定选项
        if items and len(items) >= 4:
            self.addItems(items[:4])
        else:
            # 如果没有提供足够的项目，使用默认选项
            self.addItems(["全部", "充电", "放电", "静置"])
        
        # 设置自定义委托
        self.setItemDelegate(ColoredItemDelegate(self))
        
        # 连接信号
        self.currentIndexChanged.connect(self.update_current_style)
        
        # 初始化当前项的样式
        self.update_current_style(0)
    
    def update_current_style(self, index):
        """更新当前选中项的样式"""
        if 0 <= index < 4:
            color = self.colors[index]
            self.setStyleSheet(f"""
                QComboBox {{
                    background-color: {color};
                    border: 1px solid gray;
                    padding: 1px 18px 1px 3px;
                    min-width: 6em;
                }}
                
                QComboBox::drop-down {{
                    subcontrol-origin: padding;
                    subcontrol-position: top right;
                    width: 15px;
                    border-left-width: 1px;
                    border-left-color: darkgray;
                    border-left-style: solid;
                    border-top-right-radius: 3px;
                    border-bottom-right-radius: 3px;
                }}
                
                QComboBox::down-arrow {{
                    image: url(none);
                    width: 0px;
                    height: 0px;
                }}
            """)


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("四选项四颜色 ColoredComboBox 示例")
        self.setGeometry(300, 300, 400, 200)
        
        # 创建中央部件和布局
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)
        
        # 创建第一个 ColoredComboBox，使用默认选项
        self.combo1 = ColoredComboBox()
        layout.addWidget(self.combo1)
        
        # 创建第二个 ColoredComboBox，使用自定义选项
        self.combo2 = ColoredComboBox(["低", "中", "高", "紧急"])
        layout.addWidget(self.combo2)
        
        # 创建第三个 ColoredComboBox，使用其他自定义选项
        self.combo3 = ColoredComboBox(["第一季度", "第二季度", "第三季度", "第四季度"])
        layout.addWidget(self.combo3)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())